import printer.BinaryTreeInfo;

import java.util.Comparator;
import java.util.LinkedList;
import java.util.Queue;

interface Visitor<E> {
    void visit(E element);
}

public class BinarySearchTree<E> implements BinaryTreeInfo {

    /* 根节点 */
    private Node<E> root;
    private int size;
    private Comparator<E> comparator;

    public BinarySearchTree(Comparator<E> comparator) {
        this.size = 0;
        this.comparator = comparator;
    }

    public void add(E element) {
        elementNotNullCheck(element);
        if (root == null) {
            root = new Node<>(element, null);
        } else {
            Node<E> node = root;
            Node<E> parent = node.parent;

            int compareResult = -1;
            while (node != null) {
                compareResult = compare(node.element, element);
                parent = node;
                if (compareResult > 0) {
                    node = node.left;
                } else if (compareResult < 0) {
                    node = node.right;
                } else {
                    return;
                }
            }

            Node<E> newNode = new Node<>(element, parent);
            if (compareResult > 0) {
                parent.left = newNode;
            } else if (compareResult < 0) {
                parent.right = newNode;
            } else {
                return;
            }

        }
        ++size;
    }

    public void remove(E element) {
        Node<E> nodeToRemove = node(element);
        remove(nodeToRemove);
    }

    private void remove(Node<E> node) {
        if (node == null) {
            return;
        }

        if (node.hasTwoChildren()) {
            Node<E> successor = successor(node);
            node.element = successor.element;
            node = successor;
        }

        if (node.isLeaf()) {
            /* 叶子节点 */
            if (node.isRoot()) {
                root = null;
            } else {
                if (node.parent.left == node) {
                    node.parent.left = null;
                } else if (node.parent.right == node) {
                    node.parent.right = null;
                }
            }
        } else {
            /* 非叶子节点 */
            if (node.left != null) {
                node.left.parent = node.parent;
                if (node.parent.left == node) {
                    node.parent.left = node.left;
                } else if (node.parent.right == node) {
                    node.parent.right = node.left;
                }
            } else if (node.right != null) {
                node.right.parent = node.parent;
                if (node.parent.left == node) {
                    node.parent.left = node.right;
                } else if (node.parent.right == node) {
                    node.parent.right = node.right;
                }
            }
        }
    }

    /**
     * 根据元素查找节点
     *
     * @param element 元素
     * @return 节点
     */
    private Node<E> node(E element) {
        if (root == null) {
            return null;
        }

        Node<E> node = root;

        int compareResult;
        while (node != null && (compareResult = compare(node.element, element)) != 0) {
            if (compareResult < 0) {
                node = node.right;
            } else if (compareResult > 0) {
                node = node.left;
            }
        }

        return node;
    }

    protected int compare(E e1, E e2) {
        return comparator.compare(e1, e2);
    }

    public E predecessor(E element) {
        Node<E> targetNode = node(element);
        Node<E> node = predecessor(targetNode);
        return node == null ? null : node.element;
    }

    public E successor(E element) {
        Node<E> targetNode = node(element);
        Node<E> node = successor(targetNode);
        return node == null ? null : node.element;
    }

    private Node<E> successor(Node<E> targetNode) {
        if (targetNode == null) {
            return null;
        }

        Node<E> successor;
        if (targetNode.right != null) {
            // 右子树不为空
            successor = targetNode.right;
            while (successor.left != null) {
                successor = successor.left;
            }
        } else {
            // 右子树为空
            Node<E> node = targetNode;
            successor = node.parent;

            while (node.parent != null && successor.left != node) {
                node = node.parent;
                successor = successor.parent;
            }

            return successor;
        }

        return successor;
    }

    public boolean contains(E element) {
        return node(element) != null;
    }

    /**
     * 查找指定节点的前驱节点
     *
     * @param targetNode 指定节点
     * @return 前驱节点
     */
    private Node<E> predecessor(Node<E> targetNode) {
        if (targetNode == null) {
            return null;
        }

        Node<E> predecessor;
        if (targetNode.left != null) {
            // 右子树不为空
            predecessor = targetNode.left;
            while (predecessor.right != null) {
                predecessor = predecessor.right;
            }
        } else {
            // 右子树为空
            Node<E> node = targetNode;
            predecessor = node.parent;

            while (node.parent != null && predecessor.right != node) {
                node = node.parent;
                predecessor = predecessor.parent;
            }

            return predecessor;
        }

        return predecessor;
    }

    /**
     * 节点数量
     *
     * @return 节点数量
     */
    public int size() {
        return size;
    }

    /**
     * 是否为空
     *
     * @return 是否为空
     */
    public boolean isEmpty() {
        return size() == 0;
    }

    /**
     * 清空
     */
    public void clear() {
        root = null;
        size = 0;
    }

    /**
     * 高度
     *
     * @return 高度
     */
    public int height_iter() {

        int height = 0;

        if (root == null) {
            return height;
        }

        Queue<Node<E>> queue = new LinkedList<>();
        queue.offer(root);

        int levelSize = 1;

        while (!queue.isEmpty()) {
            Node<E> node = queue.poll();
            --levelSize;

            if (node.left != null) {
                queue.offer(node.left);
            }

            if (node.right != null) {
                queue.offer(node.right);
            }

            if (levelSize == 0) {
                levelSize = queue.size();
                ++height;
            }

        }

        return height;
    }

    /* 高度 */
    public int height_recu() {
        return height_recu(root);
    }

    /* 高度 */
    private int height_recu(Node<E> node) {
        if (node == null) {
            return 0;
        }

        return 1 + Math.max(height_recu(node.left), height_recu(node.right));
    }

    /**
     * 中序遍历
     *
     * @param visitor 遍历器
     */
    public void inorderTraversal(Visitor<E> visitor) {
        inorderTraversal(root, visitor);
    }

    /**
     * 中序遍历
     *
     * @param root    根节点
     * @param visitor 遍历器
     */
    private void inorderTraversal(Node<E> root, Visitor<E> visitor) {
        if (root == null) {
            return;
        }

        inorderTraversal(root.left, visitor);
        visitor.visit(root.element);
        inorderTraversal(root.right, visitor);
    }

    /**
     * 先序便利
     *
     * @param visitor 遍历者
     */
    public void preorderTraversal(Visitor<E> visitor) {
        inorderTraversal(root, visitor);
    }

    /**
     * 先序遍历
     *
     * @param root    根节点
     * @param visitor 遍历者
     */
    private void preorderTraversal(Node<E> root, Visitor<E> visitor) {
        if (root == null) {
            return;
        }

        visitor.visit(root.element);
        preorderTraversal(root.left, visitor);
        preorderTraversal(root.right, visitor);
    }

    /**
     * 后序遍历
     *
     * @param visitor 遍历器
     */
    public void postorderTraversal(Visitor<E> visitor) {
        postorderTraversal(root, visitor);
    }

    /**
     * 后序遍历
     *
     * @param root    根节点
     * @param visitor 遍历器
     */
    private void postorderTraversal(Node<E> root, Visitor<E> visitor) {
        if (root == null) {
            return;
        }

        preorderTraversal(root.left, visitor);
        preorderTraversal(root.right, visitor);
        visitor.visit(root.element);
    }

    /**
     * 层序遍历
     *
     * @param visitor 遍历器
     */
    public void levelOrderTraversal(Visitor<E> visitor) {
        if (root == null) {
            return;
        }

        Queue<Node<E>> queue = new LinkedList<>();
        queue.offer(root);

        while (!queue.isEmpty()) {
            Node<E> node = queue.poll();
            visitor.visit(node.element);

            if (node.left != null) {
                queue.add(node.left);
            }

            if (node.right != null) {
                queue.add(node.right);
            }
        }
    }

    private void elementNotNullCheck(E element) {
        if (element == null) {
            throw new IllegalArgumentException("Element must not be null");
        }
    }

    @Override
    public Object root() {
        return root;
    }

    @Override
    public Object left(Object node) {
        return ((Node<E>) node).left;
    }

    @Override
    public Object right(Object node) {
        return ((Node<E>) node).right;
    }

    @Override
    public Object string(Object node) {
        Node<E> n = (Node<E>) node;
        return n.element + "_" + (n.parent != null ? n.parent.element : null);
    }

    private static class Node<E> {
        E element;
        Node<E> left;
        Node<E> right;
        Node<E> parent;

        public Node(E element, Node<E> parent) {
            this.element = element;
            this.parent = parent;
        }

        public boolean hasTwoChildren() {
            return left != null && right != null;
        }

        public boolean isLeaf() {
            return left == null && right == null;
        }

        public boolean isRoot() {
            return parent == null;
        }
    }
}
